#include "_2mem.h"
#include "bitmapapi.h"

#define debug kprint

/**
 * Big memory bulk management method.
 * Usually the bulk allocated by this allocator,
 * is not directlly used by applications.
 *
 * Instead the allocated bulk is used by a secondary
 * memory allocator which may be more elabrative.
 *
 * The bulk in this manager is aligned with power of 2.
 * Those discontinuous bulks are linked as a list. And the
 * head of these list is indexed by the power of the bulk size.
 *
 * All bulk has same size and is orgnized by an array.
 * The array's head pointer is stored in the allocator header.
 *
 * DY Young, @Beijing, 2024 CN
 */

typedef struct Bulk
{
   int order : 8;
   int index : 16; // max block count is 256*256
   enum
   {
      USED,
      IDLE
   } status : 2;
   int prev : 16;
   int next : 16;
} *Bulk;

typedef struct BulkAllocator
{
   word idle_map;
   struct {
      int index : 16;
   } idle_list[16];
   Bulk bulks;
   int nr_bulk;
} *BulkAllocator;

static int bit_index_of_lowest_1(word w)
{
   return __builtin_ffs(w) - 1;
}


/**
 * @brief find the first idle bulk list in the bitmap
 *        whose order is not small than specified order
 *
 * @param order the requested order
 * @param map bit map for idle list
 * @return int the list's index
 */
static int find_suitable(int order, word map)
{
   int index = -1;
   unsigned mask = 0;
   mask = ~((1 << order) - 1);
   map = map & mask;
   index = bit_index_of_lowest_1(map);
   return index;
}

static Bulk obtain_bulk_desc(int index, Bulk start)
{
   Bulk blk = NULL;
   blk = start + index;
   return blk;
}

static void link_via_bitmap(BulkAllocator ba, word bitmap, int start)
{
   int bptr = start; // bulk pointer
   int hidx = -1;    // head index of the list
   int order = -1;   // loop index poiner
   Bulk hdesc = NULL;
   Bulk ldesc = NULL;

   while (bitmap)
   {
      order = bit_index_of_lowest_1(bitmap);
      hidx = ba->idle_list[order].index;

      ldesc = obtain_bulk_desc(bptr, ba->bulks);
      ldesc->order = order;

      ba->idle_list[order].index = bptr;
      ldesc->next = hidx;

      if (hidx != -1)
      {
         hdesc = obtain_bulk_desc(hidx, ba->bulks);
         hdesc->prev = bptr;
      }
      clear_bit(&bitmap, sizeof(bitmap), order);

      debug("\t\t link bulk %d to order %d\n", bptr, order);
      bptr += (1 << order);
   }
}

/**
 * The following function is exported to other module of kernel.
 * These functions are the main API of this allocator.
 * ============================================================
 */

int allocate_bulk(BulkAllocator ba, int order)
{
   int slot = -1;
   int index = -1;
   Bulk blk = NULL;

   int expected, actual;
   word diffmap;

   expected = (1 << order);

   // set player
   slot = find_suitable(order, ba->idle_map);
   index = ba->idle_list[slot].index;

   blk = obtain_bulk_desc(index, ba->bulks);

   // update the head of the slot list
   ba->idle_list[slot].index = blk->next;

   if (blk->next == -1)
   {
      clear_bit(&ba->idle_map, sizeof(ba->idle_map), slot);
   }

   actual = (1 << slot);

   diffmap = actual - expected;

   debug("found slot %d, diffmap: %x\n", slot, diffmap);

   /**
    * Here we already found the bulk to return.
    * If the bulk is larger than the requested size,
    * we should split the builk into small blocks.
    *
    * The new generated bulks' index must  obey the
    * constrain of power of 2 rule.
    */

   ba->idle_map |= diffmap;

   link_via_bitmap(ba, diffmap, index + expected);

   blk->order = order;
   blk->status = USED;

   // release player

   return index;
}

static void remove_bulk_from_list(BulkAllocator ba, int index, int order)
{
   Bulk pblk, nblk;
   Bulk blk = obtain_bulk_desc(index, ba->bulks);
   int prev = blk->prev;
   int next = blk->next;
   debug("%d %d %d, order %d\n",prev, index, next, order);

   if(prev != -1)
   {
      pblk = obtain_bulk_desc(prev, ba->bulks);
      pblk->next = next;
   }
   if (next != -1)
   {
      nblk = obtain_bulk_desc(next, ba->bulks);
      nblk->prev = prev;
   }
   if (index == ba->idle_list[order].index)
   {
      ba->idle_list[order].index = next;
   }
   if(ba->idle_list[order].index == -1)
   {
      clear_bit(&ba->idle_map, sizeof(ba->idle_map), order);
   }
   blk->prev = -1;
   blk->next = -1;
}

void free_bulk(BulkAllocator ba, int index)
{
   Bulk blk = NULL;
   Bulk pblk = NULL;
   Bulk hblk = NULL;

   int order = 0;
   int pindex = 0;
   int hindex = index;

   blk = obtain_bulk_desc(index, ba->bulks);
   order = blk->order;

   /**
    * If current bulk's size is same with adjecent bulk's size,
    * merge them into double size bulk.
    */
   // debug("walk prev %d\n", index);
   pindex = index - (1 << order);
   while (pindex >= 0)
   {
      // walk merge to previous
      pblk = obtain_bulk_desc(pindex, ba->bulks);
      if (IDLE != pblk->status || pblk->order != order)
      {
         break;
      }

      hindex = pindex; // update the head
      // remove the bulk from the order list
      remove_bulk_from_list(ba, pindex, order);

      // merge on this condition
      // double the previous bulk's order
      order = order + 1;
      pblk->order = order;

      // walk on the previous
      pindex -= (1 << order);
   }

   // debug("walk next %d \n", index);
   pindex = hindex + (1 << order);

   hblk = obtain_bulk_desc(hindex, ba->bulks);

   while (pindex < ba->nr_bulk)
   {
      pblk = obtain_bulk_desc(pindex, ba->bulks);
      debug("watch on %d status %d\n", pindex, pblk->status);
      if (IDLE != pblk->status || pblk->order != order)
      {
         break;
      }
      // remove the bulk from the order list
      remove_bulk_from_list(ba, pindex, order);

      // merge on this condition
      // double the previous bulk's order
      order = order + 1;
      hblk->order = order;

      // walk on the previous
      pindex = hindex + (1 << order);
   }

   set_bit(&ba->idle_map, sizeof(ba->idle_map), order);

   // maintain the list
   hblk->next = ba->idle_list[order].index;
   ba->idle_list[order].index = hindex;
   hblk->status = IDLE;
   hblk->order = order;

   debug("\t\trelinked bulk %d(%d) to order %d %p\n", hindex, index, order, ba->idle_map);
}

void create_bulk_allocator(BulkAllocator ba, Bulk bulk, int nr_bulk)
{
   Bulk ptr = NULL;

   ba->bulks = bulk;
   ba->idle_map = nr_bulk;
   ba->nr_bulk = nr_bulk;

   ptr = bulk;
   for (int i = 0; i < nr_bulk; i++)
   {
      ptr->index = i;
      ptr->order = 0;
      ptr->status = IDLE;
      ptr->next = -1;
      ptr->prev = -1;
      ptr++;
   }

   for (int i = 0; i < 16; i++)
   {
      ba->idle_list[i].index = -1;
   }

   link_via_bitmap(ba, nr_bulk, 0);
}

void bulk_test()
{
#define NR 1024
#define NRT 4
   kprint("sizeof(Bulk)=%d\n", sizeof(struct Bulk));
   kprint("sizeof(struct BulkAllocator)=%d\n", sizeof(struct BulkAllocator));

   static struct Bulk bulks[NR];
   static struct BulkAllocator ba;
   create_bulk_allocator(&ba, bulks, NR);
   // allocate_bulk(&ba, 2);
   for (int i = 0; i < NRT; i++)
   {
      int idx = allocate_bulk(&ba, 0);
      kprint("allocated \033[92mbulk %d\033[0m\n\n\n", idx);
   }
   for (int i = 0; i < NRT; i++)
   {
      free_bulk(&ba, i);
   }

   for (int i = 0; i < 16; i++)
   {
      kprint("%d ", ba.idle_list[i].index);
      if(ba.idle_list[i].index != -1)
      {
         Bulk b = obtain_bulk_desc(ba.idle_list[i].index, ba.bulks);
         kprint(" next is %d\n", b->next);
      }
      else
      {
         kprint("\n");
      }
   }
   kprint("bitmap: %x\n", ba.idle_map);
}
